﻿#include "../../../includes/multimedia/rtmp/rtmpchunk.h"
#include "../../../includes/multimedia/rtmp/rtmp.h"
#include "../../../includes/QsLog/QsLog.h"
#include "../../../includes/common/common.h"

#include <random>

void WriteUint16BE(char* p, uint16_t value)
{
    p[0] = value >> 8;
    p[1] = value & 0xff;
}

void WriteUint16LE(char* p, uint16_t value)
{
    p[0] = value & 0xff;
    p[1] = value >> 8;
}

uint16_t ReadUint16BE(char* data)
{
    uint8_t* p = (uint8_t*)data;
    uint16_t value = (p[0] << 8) | p[1];
    return value;
}

uint16_t ReadUint16LE(char* data)
{
    uint8_t* p = (uint8_t*)data;
    uint16_t value = (p[1] << 8) | p[0];
    return value;
}

void WriteUint32BE(char* p, uint32_t value)
{
    p[0] = value >> 24;
    p[1] = value >> 16;
    p[2] = value >> 8;
    p[3] = value & 0xff;
}

void WriteUint32LE(char* p, uint32_t value)
{
    p[0] = value & 0xff;
    p[1] = value >> 8;
    p[2] = value >> 16;
    p[3] = value >> 24;
}

uint32_t ReadUint32BE(char* data)
{
    uint8_t* p = (uint8_t*)data;
    uint32_t value = (p[0] << 24) | (p[1] << 16) | (p[2] << 8) | p[3];
    return value;
}

uint32_t ReadUint32LE(char* data)
{
    uint8_t* p = (uint8_t*)data;
    uint32_t value = (p[3] << 24) | (p[2] << 16) | (p[1] << 8) | p[0];
    return value;
}

uint32_t ReadUint24BE(char* data)
{
    uint8_t* p = (uint8_t*)data;
    uint32_t value = (p[0] << 16) | (p[1] << 8) | p[2];
    return value;
}

uint32_t ReadUint24LE(char* data)
{
    uint8_t* p = (uint8_t*)data;
    uint32_t value = (p[2] << 16) | (p[1] << 8) | p[0];
    return value;
}

void WriteUint24BE(char* p, uint32_t value)
{
    p[0] = value >> 16;
    p[1] = value >> 8;
    p[2] = value & 0xff;
}

void WriteUint24LE(char* p, uint32_t value)
{
    p[0] = value & 0xff;
    p[1] = value >> 8;
    p[2] = value >> 16;
}

/////////////////////////////////////////////////////////////////////////////////////////

RtmpHandshake::RtmpHandshake(State state)
{
    handshake_state_ = state;
}

RtmpHandshake::~RtmpHandshake()
{

}

int RtmpHandshake::Parse(QByteArray& buffer, QByteArray& res_buf)
{
    uint8_t *buf = (uint8_t*)buffer.data();
    uint32_t buf_size = buffer.size();
    uint32_t pos = 0;
    uint32_t res_size = 0;
    std::random_device rd;

    if (handshake_state_ == HANDSHAKE_S0S1S2) {
        if (buf_size < (1 + 1536 + 1536)) { //S0S1S2
            return res_size;
        }

        if (buf[0] != RTMP_VERSION) {
            QLOG_ERROR()<< QString::asprintf("unsupported rtmp version %x", buf[0]);
            return -1;
        }

        pos += 1 + 1536 + 1536;
        res_size = 1536;
        memcpy(res_buf.data(), buf + 1, 1536); //C2
        handshake_state_ = HANDSHAKE_COMPLETE;
    }
    else if (handshake_state_ == HANDSHAKE_C0C1)
    {
        if (buf_size < 1537) { //c0c1
            return res_size;
        }
        else
        {
            if (buf[0] != RTMP_VERSION) {
                return -1;
            }

            pos += 1537;
            res_size = 1 + 1536 + 1536;
            memset(res_buf.data(), 0, 1537); //S0 S1 S2
            res_buf.data()[0] = RTMP_VERSION;

            char *p = res_buf.data(); p += 9;
            for (int i = 0; i < 1528; i++) {
                *p++ = rd();
            }
            memcpy(p, buf + 1, 1536);
            handshake_state_ = HANDSHAKE_C2;
        }
    }
    else if (handshake_state_ == HANDSHAKE_C2)
    {
        if (buf_size < 1536) { //c2
            return res_size;
        }
        else {
            pos = 1536;
            handshake_state_ = HANDSHAKE_COMPLETE;
        }
    }
    else {
        return -1;
    }

    buffer.remove(0,pos);
    return res_size;
}

int RtmpHandshake::BuildC0C1(char* buf, uint32_t buf_size)
{
    uint32_t size = 1 + 1536; //COC1
    memset(buf, 0, 1537);
    buf[0] = RTMP_VERSION;

    std::random_device rd;
    uint8_t *p = (uint8_t *)buf; p += 9;
    for (int i = 0; i < 1528; i++) {
        *p++ = rd();
    }

    return size;
}

//////////////////////////////////////////////////////////////////////////////////////

RtmpChunk::RtmpChunk()
{
    state_ = PARSE_HEADER;
    chunk_stream_id_ = -1;
    stream_id_ = kDefaultStreamId;
}

RtmpChunk::~RtmpChunk()
{

}

int RtmpChunk::Parse(QByteArray& in_buffer, RtmpMessage& out_rtmp_msg)
{
    int ret = 0;

    if (in_buffer.isEmpty()) {
        return 0;
    }

    if (state_ == PARSE_HEADER) {
        ret = ParseChunkHeader(in_buffer);
    }
    else if (state_ == PARSE_BODY)
    {
        ret = ParseChunkBody(in_buffer);
        if (ret > 0 && chunk_stream_id_ >= 0) {
            auto& rtmp_msg = rtmp_messages_[chunk_stream_id_];

            if (rtmp_msg.index == rtmp_msg.length) {
                if (rtmp_msg.timestamp >= 0xffffff) {
                    rtmp_msg._timestamp += rtmp_msg.extend_timestamp;
                }
                else {
                    rtmp_msg._timestamp += rtmp_msg.timestamp;
                }

                out_rtmp_msg = rtmp_msg;
                chunk_stream_id_ = -1;
                rtmp_msg.Clear();
            }
        }
    }

    return ret;
}

int RtmpChunk::ParseChunkHeader(QByteArray& buffer)
{
    uint32_t bytes_used = 0;
    uint8_t* buf = (uint8_t*)buffer.data();
    uint32_t buf_size = buffer.size();

    uint8_t flags = buf[bytes_used];
    bytes_used += 1;

    uint8_t csid = flags & 0x3f; // chunk stream id
    if (csid == 0) { // csid [64, 319]
        if (buf_size < (bytes_used + 2)) {
            return 0;
        }

        csid += buf[bytes_used] + 64;
        bytes_used += 1;
    }
    else if (csid == 1) { // csid [64, 65599]
        if (buf_size < (3 + bytes_used)) {
            return 0;
        }
        csid += buf[bytes_used + 1] * 255 + buf[bytes_used] + 64;
        bytes_used += 2;
    }

    uint8_t fmt = (flags >> 6); // message header type
    if (fmt >= 4) {
        return -1;
    }

    uint32_t header_len = kChunkMessageHeaderLen[fmt]; // basic_header + message_header
    if (buf_size < (header_len + bytes_used)) {
        return 0;
    }

    RtmpMessageHeader header;
    memset(&header, 0, sizeof(RtmpMessageHeader));
    memcpy(&header, buf + bytes_used, header_len);
    bytes_used += header_len;

    auto& rtmp_msg = rtmp_messages_[csid];
    chunk_stream_id_ = rtmp_msg.csid = csid;

    if (fmt == RTMP_CHUNK_TYPE_0 || fmt == RTMP_CHUNK_TYPE_1) {
        uint32_t length = ReadUint24BE((char*)header.length);
        if (rtmp_msg.length != length || !rtmp_msg.payload.isEmpty()) {
            rtmp_msg.length = length;
            rtmp_msg.payload.resize(rtmp_msg.length);
        }
        rtmp_msg.index = 0;
        rtmp_msg.type_id = header.type_id;
    }

    if (fmt == RTMP_CHUNK_TYPE_0) {
        rtmp_msg.stream_id = ReadUint24LE((char*)header.stream_id);
    }

    uint32_t timestamp = ReadUint24BE((char*)header.timestamp);
    uint32_t extend_timestamp = 0;
    if (timestamp >= 0xffffff || rtmp_msg.timestamp >= 0xffffff) {
        if (buf_size < (4 + bytes_used)) {
            return 0;
        }
        extend_timestamp = ReadUint32BE((char*)buf + bytes_used);
        bytes_used += 4;
    }

    if (rtmp_msg.index == 0) { // first chunk
        if (fmt == RTMP_CHUNK_TYPE_0) {
            // absolute timestamp
            rtmp_msg._timestamp = 0;
            rtmp_msg.timestamp = timestamp;
            rtmp_msg.extend_timestamp = extend_timestamp;
        }
        else {
            // relative timestamp (timestamp delta)
            if (rtmp_msg.timestamp >= 0xffffff) {
                rtmp_msg.extend_timestamp += extend_timestamp;
            }
            else {
                rtmp_msg.timestamp += timestamp;
            }
        }
    }

    state_ = PARSE_BODY;
    buffer.remove(0,bytes_used);
    return bytes_used;
}

int RtmpChunk::ParseChunkBody(QByteArray& buffer)
{
    uint32_t bytes_used = 0;
    uint8_t* buf = (uint8_t*)buffer.data();
    uint32_t buf_size = buffer.size();

    if (chunk_stream_id_ < 0) {
        return -1;
    }

    auto& rtmp_msg = rtmp_messages_[chunk_stream_id_];
    uint32_t chunk_size = rtmp_msg.length - rtmp_msg.index;

    if (chunk_size > in_chunk_size_) {
        chunk_size = in_chunk_size_;
    }

    if (buf_size < (chunk_size + bytes_used)) {
        return 0;
    }

    if (rtmp_msg.index + chunk_size > rtmp_msg.length) {
        return -1;
    }

    memcpy(rtmp_msg.payload.data() + rtmp_msg.index, buf + bytes_used, chunk_size);
    bytes_used += chunk_size;
    rtmp_msg.index += chunk_size;

    if (rtmp_msg.index >= rtmp_msg.length ||
        rtmp_msg.index%in_chunk_size_ == 0) {
        state_ = PARSE_HEADER;
    }

    buffer.remove(0,bytes_used);
    return bytes_used;
}

int RtmpChunk::CreateBasicHeader(uint8_t fmt, uint32_t csid, char* buf)
{
    int len = 0;

    if (csid >= 64 + 255) {
        buf[len++] = (fmt << 6) | 1;
        buf[len++] = (csid - 64) & 0xFF;
        buf[len++] = ((csid - 64) >> 8) & 0xFF;
    }
    else if (csid >= 64) {
        buf[len++] = (fmt << 6) | 0;
        buf[len++] = (csid - 64) & 0xFF;
    }
    else {
        buf[len++] = (fmt << 6) | csid;
    }
    return len;
}

int RtmpChunk::CreateMessageHeader(uint8_t fmt, RtmpMessage& rtmp_msg, char* buf)
{
    int len = 0;

    if (fmt <= 2) {
        if (rtmp_msg._timestamp < 0xffffff) {
            WriteUint24BE((char*)buf, (uint32_t)rtmp_msg._timestamp);
        }
        else {
            WriteUint24BE((char*)buf, 0xffffff);
        }
        len += 3;
    }

    if (fmt <= 1) {
        WriteUint24BE((char*)buf + len, rtmp_msg.length);
        len += 3;
        buf[len++] = rtmp_msg.type_id;
    }

    if (fmt == 0) {
        WriteUint32LE((char*)buf + len, rtmp_msg.stream_id);
        len += 4;
    }

    return len;
}

int RtmpChunk::CreateChunk(uint32_t csid, RtmpMessage& rtmp_msg, char* buf, uint32_t buf_size)
{
    uint32_t buf_offset = 0, payload_offset = 0;
    uint32_t capacity = rtmp_msg.length + rtmp_msg.length / out_chunk_size_ * 5;
    if (buf_size < capacity) {
        return -1;
    }

    buf_offset += CreateBasicHeader(0, csid, buf + buf_offset); //first chunk
    buf_offset += CreateMessageHeader(0, rtmp_msg, buf + buf_offset);
    if (rtmp_msg._timestamp >= 0xffffff) {
        WriteUint32BE((char*)buf + buf_offset, (uint32_t)rtmp_msg._timestamp);
        buf_offset += 4;
    }

    while (rtmp_msg.length > 0)
    {
        if (rtmp_msg.length > out_chunk_size_) {
            memcpy(buf + buf_offset, rtmp_msg.payload.data() + payload_offset, out_chunk_size_);
            payload_offset += out_chunk_size_;
            buf_offset += out_chunk_size_;
            rtmp_msg.length -= out_chunk_size_;

            buf_offset += CreateBasicHeader(3, csid, buf + buf_offset);
            if (rtmp_msg._timestamp >= 0xffffff) {
                WriteUint32BE(buf + buf_offset, (uint32_t)rtmp_msg._timestamp);
                buf_offset += 4;
            }
        }
        else {
            memcpy(buf + buf_offset, rtmp_msg.payload.data() + payload_offset, rtmp_msg.length);
            buf_offset += rtmp_msg.length;
            rtmp_msg.length = 0;
            break;
        }
    }

    return buf_offset;
}
